#!/usr/bin/env python3
"""
Advanced Go Binary Analysis for Mapping Table Extraction
"""

import struct
import os
import re

def analyze_go_pclntab(binary_path):
    """Analyze Go pclntab (program counter line table) to find function addresses"""
    
    print("=== Analyzing Go .gopclntab section ===")
    
    with open(binary_path, 'rb') as f:
        data = f.read()
    
    # Look for Go pclntab header (magic numbers)
    pclntab_magic = [
        b'\xfb\xff\xff\xff\x00\x00',  # Go 1.2+
        b'\xfa\xff\xff\xff\x00\x00',  # Go 1.16+
        b'\xf0\xff\xff\xff\x00\x00'   # Go 1.18+
    ]
    
    pclntab_start = None
    for magic in pclntab_magic:
        pos = data.find(magic)
        if pos != -1:
            pclntab_start = pos
            print(f"Found Go pclntab header at 0x{pos:x}")
            break
    
    if not pclntab_start:
        print("No Go pclntab header found")
        return None
    
    # Parse the pclntab structure
    try:
        # Read header
        header = struct.unpack('<IIHH', data[pclntab_start:pclntab_start+12])
        magic, pad, minlc, ptrsize = header
        print(f"pclntab: magic=0x{magic:x}, minlc={minlc}, ptrsize={ptrsize}")
        
        # Calculate function table location
        funcnametab_offset = pclntab_start + 8 + ptrsize
        functab_offset = funcnametab_offset + 4
        
        # Read number of functions
        nfunc = struct.unpack('<I', data[functab_offset:functab_offset+4])[0]
        print(f"Number of functions: {nfunc}")
        
        # Look for encryption functions
        encrypt_functions = []
        func_table_start = functab_offset + 4
        
        for i in range(min(nfunc, 1000)):  # Limit to avoid huge output
            func_entry = func_table_start + i * (2 * ptrsize)
            if func_entry + 2 * ptrsize > len(data):
                break
                
            if ptrsize == 8:
                pc, funcoff = struct.unpack('<QQ', data[func_entry:func_entry+16])
            else:
                pc, funcoff = struct.unpack('<II', data[func_entry:func_entry+8])
            
            # Get function name
            name_offset = pclntab_start + funcoff + 8  # Skip pc and name length
            if name_offset < len(data):
                # Read name length
                name_len = struct.unpack('<I', data[pclntab_start + funcoff + 4:pclntab_start + funcoff + 8])[0]
                if name_len < 1000 and name_offset + name_len < len(data):
                    func_name = data[name_offset:name_offset + name_len].decode('utf-8', errors='ignore')
                    
                    if 'encrypt' in func_name or 'cosy/util' in func_name:
                        encrypt_functions.append((pc, func_name))
                        print(f"Found function: {func_name} at PC 0x{pc:x}")
        
        return encrypt_functions
        
    except Exception as e:
        print(f"Error parsing pclntab: {e}")
        return None

def search_near_functions(binary_path, function_addresses):
    """Search for data structures near function addresses"""
    
    print("\n=== Searching near function addresses ===")
    
    with open(binary_path, 'rb') as f:
        data = f.read()
    
    # Known mapping
    known_mappings = {
        0xa1: 0x6d,  # m
        0xf3: 0x6f,  # o  
        0x6b: 0x64,  # d
        0x31: 0x65,  # e
        0xb0: 0x6c,  # l
        0x69: 0x22,  # "
        0x01: 0x3a,  # :
        0x90: 0x20   # space
    }
    
    candidates = []
    
    for pc, func_name in function_addresses:
        if 'encrypt' in func_name.lower():
            print(f"\nSearching around {func_name} (PC: 0x{pc:x})")
            
            # Convert PC to file offset (approximate)
            # Go binaries typically have text section starting at 0x400000
            file_offset = pc - 0x400000
            if file_offset < 0 or file_offset >= len(data):
                print(f"PC 0x{pc:x} out of file range")
                continue
            
            # Search in a range around the function
            search_start = max(0, file_offset - 0x10000)  # 64KB before
            search_end = min(len(data), file_offset + 0x10000)  # 64KB after
            
            # Look for 256-byte arrays with our known patterns
            for offset in range(search_start, search_end - 256, 16):
                segment = data[offset:offset + 256]
                
                # Check known mappings
                matches = 0
                for src, dst in known_mappings.items():
                    if src < 256 and segment[src] == dst:
                        matches += 1
                
                if matches >= 4:  # Looser criteria
                    print(f"  Found candidate at 0x{offset:x} ({matches}/8 matches)")
                    candidates.append((offset, segment, matches, func_name))
    
    return candidates

def search_for_base64_alphabet(binary_path):
    """Search for the custom Base64 alphabet in various forms"""
    
    print("\n=== Advanced Base64 Alphabet Search ===")
    
    with open(binary_path, 'rb') as f:
        data = f.read()
    
    # Try different representations of the custom alphabet
    custom_base64 = "!#%&()*,.@ABCDEFGHIJKLMNOPQRSTUVWXYZ^_abcdefghijklmnopqrstuvwxyz"
    
    # 1. Direct byte string
    alphabet_bytes = custom_base64.encode('ascii')
    pos = data.find(alphabet_bytes)
    if pos != -1:
        print(f"✅ Found complete alphabet at 0x{pos:x}")
        return pos
    
    # 2. Look for individual segments
    segments = [
        "!#%&()*,.@",
        "ABCDEFGHIJKLMNOPQRSTUVWXYZ", 
        "^_abcdefghijklmnopqrstuvwxyz"
    ]
    
    segment_positions = []
    for segment in segments:
        pos = data.find(segment.encode('ascii'))
        if pos != -1:
            segment_positions.append((pos, segment))
            print(f"Found segment '{segment}' at 0x{pos:x}")
    
    # 3. Look for the alphabet stored as a lookup table
    # The encoder might use it as a 64-byte lookup array
    if len(segment_positions) >= 2:
        print("Multiple segments found - alphabet might be stored separately")
    
    # 4. Search for the alphabet in reverse (decoder table)
    # Create reverse mapping
    reverse_map = {}
    for i, char in enumerate(custom_base64):
        reverse_map[ord(char)] = i
    
    # Look for patterns where ASCII values map to positions
    print("Searching for reverse lookup patterns...")
    for offset in range(0, len(data) - 128, 8):
        segment = data[offset:offset + 128]
        
        # Check if this could be a reverse lookup table
        matches = 0
        for char_pos, expected_idx in reverse_map.items():
            if char_pos < 128 and segment[char_pos] == expected_idx:
                matches += 1
        
        if matches >= 10:  # At least 10 characters match
            print(f"Found reverse lookup candidate at 0x{offset:x} ({matches} matches)")
            
            # Save this candidate
            with open(f'reverse_lookup_0x{offset:x}.bin', 'wb') as f:
                f.write(segment)
    
    return None

def create_runtime_extractor():
    """Create a script to extract mapping table at runtime"""
    
    script_content = '''#!/usr/bin/env python3
"""
Runtime mapping table extractor using memory dumps
Run this while Lingma is active
"""

import subprocess
import re
import struct

def find_lingma_process():
    """Find running Lingma process"""
    try:
        result = subprocess.run(['pgrep', '-f', 'Lingma'], 
                               capture_output=True, text=True)
        if result.returncode == 0:
            return result.stdout.strip().split()[0]
    except:
        pass
    return None

def dump_process_memory(pid):
    """Dump process memory using gcore"""
    try:
        subprocess.run(['sudo', 'gcore', pid], check=True)
        return f'core.{pid}'
    except subprocess.CalledProcessError:
        print(f"Failed to dump memory for PID {pid}")
        return None

def search_memory_dump(core_file):
    """Search memory dump for mapping table"""
    known_mappings = {
        0xa1: 0x6d,  # m
        0xf3: 0x6f,  # o  
        0x6b: 0x64,  # d
        0x31: 0x65,  # e
        0xb0: 0x6c,  # l
        0x69: 0x22,  # "
        0x01: 0x3a,  # :
        0x90: 0x20   # space
    }
    
    print(f"Searching memory dump {core_file}...")
    with open(core_file, 'rb') as f:
        data = f.read()
    
    candidates = []
    for offset in range(0, len(data) - 256, 8):
        segment = data[offset:offset + 256]
        
        matches = 0
        for src, dst in known_mappings.items():
            if segment[src] == dst:
                matches += 1
        
        if matches >= 6:
            print(f"Found mapping table at offset 0x{offset:x} ({matches}/8 matches)")
            with open(f'runtime_mapping_0x{offset:x}.bin', 'wb') as f:
                f.write(segment)
            candidates.append((offset, segment))
    
    return candidates

if __name__ == "__main__":
    print("Runtime Mapping Table Extractor")
    print("================================")
    
    pid = find_lingma_process()
    if not pid:
        print("❌ Lingma process not found")
        print("Please start Lingma and try again")
        exit(1)
    
    print(f"Found Lingma process: PID {pid}")
    
    core_file = dump_process_memory(pid)
    if not core_file:
        exit(1)
    
    candidates = search_memory_dump(core_file)
    if candidates:
        print(f"✅ Found {len(candidates)} mapping table candidates")
    else:
        print("❌ No mapping tables found in memory dump")
'''
    
    with open('runtime_extractor.py', 'w') as f:
        f.write(script_content)
    
    print("Created runtime_extractor.py - run this when Lingma is active")

if __name__ == "__main__":
    binary_path = "Lingma"
    
    if not os.path.exists(binary_path):
        print(f"Binary {binary_path} not found!")
        exit(1)
    
    print("Advanced Go Binary Analysis for Tongyi Lingma")
    print("=" * 50)
    
    # 1. Analyze pclntab to find function addresses
    encrypt_functions = analyze_go_pclntab(binary_path)
    
    if encrypt_functions:
        # 2. Search near function addresses for data
        candidates = search_near_functions(binary_path, encrypt_functions)
        
        if candidates:
            print(f"\n✅ Found {len(candidates)} mapping table candidates near functions")
            for i, (offset, segment, matches, func_name) in enumerate(candidates[:3]):
                filename = f"function_mapping_{i+1}_0x{offset:x}.bin"
                with open(filename, 'wb') as f:
                    f.write(segment)
                print(f"Saved candidate {i+1} to {filename} (from {func_name})")
    
    # 3. Search for Base64 alphabet
    search_for_base64_alphabet(binary_path)
    
    # 4. Create runtime extractor
    create_runtime_extractor()
    
    print("\n" + "=" * 50)
    print("Analysis complete!")
    print("\nNext steps:")
    print("1. Check the generated mapping candidates")
    print("2. Start Lingma and run runtime_extractor.py")
    print("3. Use GDB/strace to trace actual function calls")